{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating a custom imputer in scikit-learn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd \n",
    "import numpy as np\n",
    "from sklearn.base import BaseEstimator, TransformerMixin\n",
    "from sklearn.utils.validation import check_is_fitted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The scikit-learn version is 0.22.2.post1.\n"
     ]
    }
   ],
   "source": [
    "import sklearn \n",
    "print(f'The scikit-learn version is {sklearn.__version__}.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>sample_name</th>\n",
       "      <th>variant</th>\n",
       "      <th>height</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>B</td>\n",
       "      <td>a</td>\n",
       "      <td>171.678012</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>B</td>\n",
       "      <td>a</td>\n",
       "      <td>166.292437</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>A</td>\n",
       "      <td>a</td>\n",
       "      <td>168.286427</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>B</td>\n",
       "      <td>a</td>\n",
       "      <td>158.773399</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>B</td>\n",
       "      <td>a</td>\n",
       "      <td>155.756804</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  sample_name variant      height\n",
       "0           B       a  171.678012\n",
       "1           B       a  166.292437\n",
       "2           A       a  168.286427\n",
       "3           B       a  158.773399\n",
       "4           B       a  155.756804"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# setting the seed\n",
    "np.random.seed(42)\n",
    "\n",
    "# generating the two samples\n",
    "sample_a = pd.DataFrame({'sample_name': 'A',\n",
    "                         'variant': np.random.choice(['a', 'b'], size=100),\n",
    "                         'height': np.random.normal(loc=170, scale=5, size=100)})\n",
    "\n",
    "sample_b = pd.DataFrame({'sample_name': 'B',\n",
    "                         'variant': 'a',\n",
    "                         'height': np.random.normal(loc=165, scale=7, size=100)})\n",
    "\n",
    "# concatenating the samples\n",
    "df = pd.concat([sample_a, sample_b], axis=0, ignore_index=True) \\\n",
    "       .sample(frac=1) \\\n",
    "       .reset_index(drop=True)\n",
    "\n",
    "# preview the data\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "ind_to_replace = np.random.choice(range(len(df)), 10, replace=False)\n",
    "df.loc[ind_to_replace, 'height'] = np.nan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>sample_name</th>\n",
       "      <th>variant</th>\n",
       "      <th>height</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>B</td>\n",
       "      <td>a</td>\n",
       "      <td>171.678012</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>B</td>\n",
       "      <td>a</td>\n",
       "      <td>166.292437</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>A</td>\n",
       "      <td>a</td>\n",
       "      <td>168.286427</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>B</td>\n",
       "      <td>a</td>\n",
       "      <td>158.773399</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>B</td>\n",
       "      <td>a</td>\n",
       "      <td>155.756804</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>195</th>\n",
       "      <td>B</td>\n",
       "      <td>a</td>\n",
       "      <td>178.060422</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>196</th>\n",
       "      <td>A</td>\n",
       "      <td>b</td>\n",
       "      <td>167.482622</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>197</th>\n",
       "      <td>B</td>\n",
       "      <td>a</td>\n",
       "      <td>169.803821</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>198</th>\n",
       "      <td>A</td>\n",
       "      <td>a</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>199</th>\n",
       "      <td>A</td>\n",
       "      <td>b</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>200 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "    sample_name variant      height\n",
       "0             B       a  171.678012\n",
       "1             B       a  166.292437\n",
       "2             A       a  168.286427\n",
       "3             B       a  158.773399\n",
       "4             B       a  155.756804\n",
       "..          ...     ...         ...\n",
       "195           B       a  178.060422\n",
       "196           A       b  167.482622\n",
       "197           B       a  169.803821\n",
       "198           A       a         NaN\n",
       "199           A       b         NaN\n",
       "\n",
       "[200 rows x 3 columns]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating the custom imputer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GroupImputer(BaseEstimator, TransformerMixin):\n",
    "    '''\n",
    "    Class used for imputing missing values in a pd.DataFrame using either mean or median of a group.\n",
    "    \n",
    "    Parameters\n",
    "    ----------    \n",
    "    group_cols : list\n",
    "        List of columns used for calculating the aggregated value \n",
    "    target : str\n",
    "        The name of the column to impute\n",
    "    metric : str\n",
    "        The metric to be used for remplacement, can be one of ['mean', 'median']\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    X : array-like\n",
    "        The array with imputed values in the target column\n",
    "    '''\n",
    "    def __init__(self, group_cols, target, metric='mean'):\n",
    "        \n",
    "        assert metric in ['mean', 'median'], 'Unrecognized value for metric, should be mean/median'\n",
    "        assert type(group_cols) == list, 'group_cols should be a list of columns'\n",
    "        assert type(target) == str, 'target should be a string'\n",
    "        \n",
    "        self.group_cols = group_cols\n",
    "        self.target = target\n",
    "        self.metric = metric\n",
    "    \n",
    "    def fit(self, X, y=None):\n",
    "        \n",
    "        assert pd.isnull(X[self.group_cols]).any(axis=None) == False, 'There are missing values in group_cols'\n",
    "        \n",
    "        impute_map = X.groupby(self.group_cols)[self.target].agg(self.metric) \\\n",
    "                                                            .reset_index(drop=False)\n",
    "        \n",
    "        self.impute_map_ = impute_map\n",
    "        \n",
    "        return self \n",
    "    \n",
    "    def transform(self, X, y=None):\n",
    "        \n",
    "        # make sure that the imputer was fitted\n",
    "        check_is_fitted(self, 'impute_map_')\n",
    "        \n",
    "        X = X.copy()\n",
    "        \n",
    "        for index, row in self.impute_map_.iterrows():\n",
    "            ind = (X[self.group_cols] == row[self.group_cols]).all(axis=1)\n",
    "            X.loc[ind] = X.loc[ind].fillna(row[self.target])\n",
    "        \n",
    "        return X.values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "df contains 10 missing values.\n",
      "df_imp contains 0 missing values.\n"
     ]
    }
   ],
   "source": [
    "imp = GroupImputer(group_cols=['sample_name', 'variant'], \n",
    "                   target='height', \n",
    "                   metric='mean')\n",
    "\n",
    "df_imp = pd.DataFrame(imp.fit_transform(df), \n",
    "                      columns=df.columns)\n",
    "\n",
    "print(f'df contains {sum(pd.isnull(df.height))} missing values.')\n",
    "print(f'df_imp contains {sum(pd.isnull(df_imp.height))} missing values.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>sample_name</th>\n",
       "      <th>variant</th>\n",
       "      <th>height</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>B</td>\n",
       "      <td>a</td>\n",
       "      <td>171.678</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>B</td>\n",
       "      <td>a</td>\n",
       "      <td>166.292</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>A</td>\n",
       "      <td>a</td>\n",
       "      <td>168.286</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>B</td>\n",
       "      <td>a</td>\n",
       "      <td>158.773</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>B</td>\n",
       "      <td>a</td>\n",
       "      <td>155.757</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>195</th>\n",
       "      <td>B</td>\n",
       "      <td>a</td>\n",
       "      <td>178.06</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>196</th>\n",
       "      <td>A</td>\n",
       "      <td>b</td>\n",
       "      <td>167.483</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>197</th>\n",
       "      <td>B</td>\n",
       "      <td>a</td>\n",
       "      <td>169.804</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>198</th>\n",
       "      <td>A</td>\n",
       "      <td>a</td>\n",
       "      <td>170.362</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>199</th>\n",
       "      <td>A</td>\n",
       "      <td>b</td>\n",
       "      <td>169.51</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>200 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "    sample_name variant   height\n",
       "0             B       a  171.678\n",
       "1             B       a  166.292\n",
       "2             A       a  168.286\n",
       "3             B       a  158.773\n",
       "4             B       a  155.757\n",
       "..          ...     ...      ...\n",
       "195           B       a   178.06\n",
       "196           A       b  167.483\n",
       "197           B       a  169.804\n",
       "198           A       a  170.362\n",
       "199           A       b   169.51\n",
       "\n",
       "[200 rows x 3 columns]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_imp"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
